import numpy as np
import platgo as pg


class DTLZ9(pg.Problem):

    def __init__(self, M: int = 3) -> None:
        self.name = 'DTLZ9'
        self.type['multi'], self.type['many'], self.type['real'], self.type['large'], self.type['expensive'], self.type['constrained'] = [
                                                                                                                   True] * 6
        self.M = M
        self.D = M * 10
        lb = [0] * self.D
        ub = [1] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        decs = pop.decs ** 0.1
        pop.objv = np.zeros((pop.N, self.M))
        pop.cv = np.zeros((pop.N, self.M))
        for i in range(1, self.M + 1):
            f = np.sum(decs[:, ((i - 1) * 10):i * 10], axis=1, keepdims=True) * 0.1
            pop.objv[:, i - 1:i] = f[:, :]

        pop.cv = 1 - pop.objv[:, :(self.M - 1)] ** 2 - pop.objv[:, -1:] ** 2

    def get_optimal(self) -> np.ndarray:
        # 目标空间均匀采样函数还未完成
        raise NotImplementedError("get optimal has not been implemented")


if __name__ == '__main__':
    d = DTLZ9()
    # pop = pg.Population(decs=np.random.uniform(0, 1, (10, 30)))
    pop = d.init_pop(5)
    print(d.borders)
    print(pop.decs)
    d.cal_obj(pop)  # 计算目标函数值
    print(pop.objv)
    print(pop.cv)
